import math

import numpy as np
from matplotlib import pyplot as plt

from BasisFunctions import BasisFuction
from Curve import Curve


def calcularG_bar(n):
    I_2 = np.identity(2)

    s = np.array([(i - 1) / n for i in range(1, n + 1)])
    h = math.ceil((n - 1) / 2) if math.ceil((n - 1) / 2) > 6 else 6
    while True:
        basic = BasisFuction(n, h, s)
        G_i = basic.construct()

        G_bar = np.vstack([G_i[i] for i in range(n)])
        rank = np.linalg.matrix_rank(G_bar)
        if rank == 2 * n:
            break
        h += 1
        del basic
    return h, G_i, G_bar


def R(theta, l):
    return np.array(
        [[np.cos(theta), -l * np.sin(theta)], [np.sin(theta), l * np.cos(theta)]]
    )


def draw(curve: Curve, x, y, ax):
    a = np.linspace(0, 1, 1000)
    sample_x = curve.fx(a)
    sample_y = curve.fy(a)

    # 清除当前绘图内容
    ax.cla()

    # 绘制曲线
    ax.plot(sample_x, sample_y, "b-", label="Curve")

    # 添加采样点标记（红色圆圈+星号边框）
    ax.scatter(
        x[0],
        y[0],
        marker="o",  # 圆形标记
        c="yellow",  # 填充颜色
        linewidths=2,  # 边框粗细
        label="Samples",
    )

    ax.scatter(
        x[1:],
        y[1:],
        marker="o",  # 圆形标记
        c="red",  # 填充颜色
        linewidths=2,  # 边框粗细
        label="Samples",
    )
    ax.axis("equal")
    ax.legend()
    plt.draw()  # 更新绘图内容
    plt.pause(1)  # 暂停0.1秒


if __name__ == "__main__":

    fig, ax = plt.subplots(figsize=(8, 8))
    n = 1000
    l = 0.01
    k1 = 1
    k2 = 1
    d = np.array([0.1, 0.1]).reshape(2, 1)
    # a = np.array(
    #     [
    #         [1, 1, 1, 0, 0, 0],
    #         [1, 1, 0, 1, 0, 0],
    #         [1, 0, 1, 0, 1, 1],
    #         [0, 1, 0, 1, 0, 0],
    #         [0, 0, 1, 0, 1, 0],
    #         [0, 0, 1, 0, 0, 1],
    #     ]
    # )
    a = np.array(
        [
            [1, 1, 1, 0, 0, 0],
            [1, 1, 0, 0, 0, 0],
            [1, 0, 1, 0, 0, 0],
            [0, 1, 0, 1, 0, 0],
            [0, 0, 1, 0, 1, 0],
            [0, 0, 1, 0, 0, 1],
        ]
    )

    u = np.zeros((n, 2, 1))
    x = np.zeros((n, 2, 1))
    r = np.zeros((n, 2, 2))
    delta = np.zeros((n, 2, 1))
    kxie_norm = []

    curve = Curve()

    h, G_i, G_bar = calcularG_bar(n)
    kxi = curve.estima(h)

    G_plus = np.linalg.pinv(G_bar)

    agent_x = np.random.uniform(low=0, high=10, size=(n,))
    agent_y = np.random.uniform(low=0, high=10, size=(n,))
    agent_theta = np.random.uniform(low=0, high=2 * np.pi, size=(n,))

    for t in range(100):
        x_bar = np.vstack([agent_x[0], agent_y[0]])
        for i in range(1, n):
            x_bar = np.vstack([x_bar, agent_x[i], agent_y[i]])

        kxi_e = G_plus @ x_bar - kxi
        kxie_norm.append(np.linalg.norm(kxi_e))
        for i in range(n):
            x[i] = np.vstack([agent_x[i], agent_y[i]])
            # print(x[i].shape)
            r[i] = R(agent_theta[i], l)
            delta[i] = k2 * r[i].T @ (x[i] - G_i[i] @ kxi)

            # if i == 0:
            #     u[i] = -k1 * (x[i] - G_i[i] @ kxi) - k2 * r[i] @ delta[i]
            # else:
            #     neiber_sum = np.zeros((2, 2 * (2 * h + 1)))
            #     for j in range(n):
            #         if a[i][j] == 1:
            #             neiber_sum += a[i][j] * (G_i[i] - G_i[j])
            #     u[i] = -k1 * neiber_sum @ kxi_e - k2 * r[i] @ delta[i]
            u[i] = -k1 * (x[i] - G_i[i] @ kxi) - k2 * r[i] @ delta[i]

            x[i] = u[i] + r[i] @ d
            agent_x[i] += x[i][0][0]
            agent_y[i] += x[i][1][0]
            agent_theta[i] += u[i][1][0] + d[1][0]
            # break
        draw(curve, agent_x, agent_y, ax)
        # break
    fig2, ax2 = plt.subplots(figsize=(8, 8))  # 创建第二个窗口
    ax2.plot(kxie_norm, color='b', linestyle='--', label='cos(x)')
    ax2.set_title("kxi_e_norm")
    ax2.legend()
    plt.show()
